package com.squirrel.michale;

import java.util.*;

public class Day20240428_LeetCode003123 {

    public boolean[] findAnswer(int n, int[][] edges) {

        List<int[]>[] graph = new ArrayList[n];
        Arrays.setAll(graph, i -> new ArrayList<>());
        for (int i = 0; i < edges.length; i++) {
            int[] e = edges[i];
            int x = e[0], y = e[1], w = e[2];
            graph[x].add(new int[]{y, w, i});
            graph[y].add(new int[]{x, w, i});
        }

        long[] dis = new long[n];
        Arrays.fill(dis, Long.MAX_VALUE);
        dis[0] = 0;
        PriorityQueue<long[]> pq = new PriorityQueue<>((a, b) -> Long.compare(a[0], b[0]));
        pq.offer(new long[]{0, 0});
        while (!pq.isEmpty()) {
            long[] dxPair = pq.poll();
            long dx = dxPair[0];
            int x = (int) dxPair[1];
            if (dx > dis[x]) {
                continue;
            }
            for (int[] t : graph[x]) {
                int y = t[0];
                int w = t[1];
                long newDis = dx + w;
                if (newDis < dis[y]) {
                    dis[y] = newDis;
                    pq.offer(new long[]{newDis, y});
                }
            }
        }

        boolean[] ans = new boolean[edges.length];
        // 图不连通
        if (dis[n - 1] == Long.MAX_VALUE) {
            return ans;
        }

        // 从终点出发 BFS
        boolean[] vis = new boolean[n];
        vis[n - 1] = true;
        Queue<Integer> q = new ArrayDeque<>();
        q.offer(n - 1);
        while (!q.isEmpty()) {
            int y = q.poll();
            for (int[] t : graph[y]) {
                int x = t[0];
                int w = t[1];
                int i = t[2];
                if (dis[x] + w != dis[y]) {
                    continue;
                }
                ans[i] = true;
                if (!vis[x]) {
                    vis[x] = true;
                    q.offer(x);
                }
            }
        }
        return ans;


    }


    public static void main(String[] args) {
        //[[0,1,4],[0,2,1],[1,3,2],[1,4,3],[1,5,1],[2,3,1],[3,5,3],[4,5,2]]

        int[][] edges = new int[][]{{0, 1, 4}, {0, 2, 1}, {1, 3, 2}, {1, 4, 3}, {1, 5, 1}, {2, 3, 1}, {3, 5, 3}, {4, 5, 2}};
        int n = 6;


        Day20240428_LeetCode003123 leetCode003123 = new Day20240428_LeetCode003123();
        boolean[] answer = leetCode003123.findAnswer(n, edges);
        System.out.println(Arrays.toString(answer));


    }

}
